import argparse
import os
import random
from collections import namedtuple

import numpy as np
from sklearn.metrics import roc_auc_score

num_sites = 100000000
def parse_args():
    parser = argparse.ArgumentParser("")
    parser.add_argument("--input",'-i', type=str,required=True)
    parser.add_argument('--result_file', type=str, required=False,
                        help='the result file are going to save')
    parser.add_argument('--motif', type=str, required=False,
                        help='the result file are going to save')
    return parser.parse_args()

def evaluate(pred_p,pred_n,result_file):
    pr_writer = open(result_file, 'w')
    pr_writer.write("tested_type\tTP\tFN\tTN\tFP\t"
                    "accuracy\trecall\tspecificity\tprecision\t"
                    "fallout\tmiss_rate\tFDR\tNPV\tAUC\tsamplenum\n")
    metrics = []
    num_rounds=5
    if num_sites >= len(pred_p) and num_sites >= len(pred_n):
        num_rounds = 1
    for i in range(num_rounds):
        random.shuffle(pred_p)
        random.shuffle(pred_n)
        tp = 0
        fp = 0
        tn = 0
        fn = 0
        pred_pos=pred_p[:num_sites]
        pred_neg=pred_n[:num_sites]
        called = 0
        correct = 0
        y_truelabel = []
        y_scores = []
        for pred in pred_pos:
            if pred>0.5:
                tp+=1
            else:
                fn+=1
            y_truelabel.append(1)
            y_scores.append(pred)
        for pred in pred_neg:
            if pred>0.5:
                fp+=1
            else:
                tn+=1
            y_truelabel.append(0)
            y_scores.append(pred)
        precision, recall, specificity, accuracy = 0, 0, 0, 0
        fall_out, miss_rate, fdr, npv, = 0, 0, 0, 0
        auroc = 0
        if num_sites > 0:
            accuracy = float(tp + tn) / (len(pred_neg)+len(pred_pos))
            if tp + fp > 0:
                precision = float(tp) / (tp + fp)
                fdr = float(fp) / (tp + fp)  # false discovery rate
            else:
                precision = 0
                fdr = 0
            if tp + fn > 0:
                recall = float(tp) / (tp + fn)
                miss_rate = float(fn) / (tp + fn)  # false negative rate
            else:
                recall = 0
                miss_rate = 0
            if tn + fp > 0:
                specificity = float(tn) / (tn + fp)
                fall_out = float(fp) / (fp + tn)  # false positive rate
            else:
                specificity = 0
                fall_out = 0
            if tn + fn > 0:
                npv = float(tn) / (tn + fn)  # negative predictive value
            else:
                npv = 0
            auroc = roc_auc_score(np.array(y_truelabel), np.array(y_scores))
        metrics.append([tp, fn, tn, fp, accuracy, recall, specificity, precision,
                                    fall_out, miss_rate, fdr, npv, auroc, num_sites*2])
    metrics = np.array(metrics, dtype=float)

    metrics_mean = np.mean(metrics, 0)
    mean_tpfntnfp = "\t".join([str(round(x, 1)) for x in metrics_mean[:4]])
    mean_perf = "\t".join([str(round(x, 4)) for x in metrics_mean[4:13]])
    mean_numlen = str(round(metrics_mean[13]))
    pr_writer.write("\t".join([str(num_sites), mean_tpfntnfp, mean_perf, mean_numlen,
                                ]) + "\n")

    metrics_std = np.std(metrics, 0)
    std_tpfntnfp = "\t".join([str(round(x, 1)) for x in metrics_std[:4]])
    std_perf = "\t".join([str(round(x, 4)) for x in metrics_std[4:13]])
    std_numlen = str(round(metrics_std[13]))
    pr_writer.write("\t".join([str(num_sites) + "_std", std_tpfntnfp, std_perf, std_numlen,
                                ]) + "\n")
    pr_writer.flush()
    pr_writer.close()

def get_refloc_of_methysite_in_motif(seqstr, motifset, methyloc_in_motif=0):
    """

    :param seqstr:
    :param motifset:
    :param methyloc_in_motif: 0-based
    :return:
    """
    motifset = set(motifset)
    strlen = len(seqstr)
    motiflen = len(list(motifset)[0])
    sites = []
    for i in range(0, strlen - motiflen + 1):
        if seqstr[i : i + motiflen] in motifset:
            sites.append(i + methyloc_in_motif)
    return sites

def read_feature(filename,result_file,motif):
    rf = open(filename)
    pred_deepsignals_p=[]
    pred_dorados_p=[]
    pred_means_p=[]
    pred_deepsignals_n=[]
    pred_dorados_n=[]
    pred_means_n=[]
    #pos_pos=[]

    for line in rf:
        words = line.strip().split("\t")
        pred_deepsignal = [float(x)for x in words[7].split(',')]
        pred_dorado =[float(x)for x in words[6].split(',')]
        pred_mean=[float(x)for x in words[8].split(',')]
        bis_label=[float(x)for x in words[15].split(',')]
        seq=words[3]
        pred_pos=[int(x)for x in words[5].split(',')]
        if motif is not None:
            sites=get_refloc_of_methysite_in_motif(seq,motif)
        for i in range(len(bis_label)):
            # if motif is not None:
            #     if motif=='CHG':
            #         if len(seq)<=pred_pos[i]+2:
            #             continue
            #         if seq[pred_pos[i]+1]=='G' :
            #             continue
            #         if seq[pred_pos[i]+2]!='G' :
            #             continue
            #     elif motif=='CHH':
            #         if len(seq)<=pred_pos[i]+2:
            #             continue
            #         if seq[pred_pos[i]+1]=='G' or seq[pred_pos[i]+2]=='G' :
            #             continue
            if motif is not None:
                if pred_pos[i] not in sites:
                    continue            
            if float(bis_label[i])==100.0:
                if motif is None:
                    if pred_deepsignal[i]==-1 or pred_dorado[i]==-1:
                        continue
                else:
                    if pred_deepsignal[i]==-1:
                        continue
                pred_deepsignals_p.append(pred_deepsignal[i])
                if motif is None:
                    pred_dorados_p.append(pred_dorado[i])
                    pred_means_p.append(pred_mean[i])
            elif float(bis_label[i])==0.0:
                if motif is None:
                    if pred_deepsignal[i]==-1 or pred_dorado[i]==-1:
                        continue
                else:
                    if pred_deepsignal[i]==-1:
                        continue
                pred_deepsignals_n.append(pred_deepsignal[i])
                if motif is None:
                    pred_dorados_n.append(pred_dorado[i])
                    pred_means_n.append(pred_mean[i])
    # if len(pred_deepsignals_p)<=num_sites or len(pred_deepsignals_n)<=num_sites:
    #     print('error')
    #     print(len(pred_deepsignals_p))
    #     print(len(pred_deepsignals_n))
    #     exit()
    evaluate(pred_deepsignals_p,pred_deepsignals_n,result_file+'deepsignal.txt')
    if motif is None:
        evaluate(pred_dorados_p,pred_dorados_n,result_file+'dorado.txt')
        
        evaluate(pred_means_p,pred_means_n,result_file+'mean.txt')
    
    #return pred_deepsignals_p,pred_dorados_p,pred_means_p,pred_deepsignals_n,pred_dorados_n,pred_means_n

if __name__ == '__main__':
    args=parse_args()
    read_feature(args.input,args.result_file,args.motif)